/*
 * armavec.h
 *
 *  Created on: Mar 6, 2011
 *      Author: tqlong
 */

#ifndef ARMAVEC_H_
#define ARMAVEC_H_

#include <armadillo>
/** Interface to arma::vec
 * Adding dim(), get(), ref() to access element
 */
class ArmaVec : public arma::vec
{
  typedef arma::vec __Base;
public:
  ArmaVec(int d) : __Base(d) {}
  ArmaVec() {}
  ArmaVec(const __Base& v): __Base(v) {}
//  ArmaVec& operator=(const ArmaVec&v) { __Base::operator=(v); return *this; }
  template<typename T1, typename T2, typename T3>
  ArmaVec& operator=(const arma::eGlue<T1,T2,T3> &v) { __Base::operator=(v); return *this; }
  template<typename T1, typename T2, typename T3>
  ArmaVec(const arma::eGlue<T1,T2,T3> &v) : __Base(v) {}
  template<typename T1, typename T2>
  ArmaVec(const arma::eOp<T1,T2> &v) : __Base(v) {}
  double get(int i) const { return at(i); }
  double& ref(int i) { return at(i); }
  int dim() const { return __Base::n_elem; }
  ArmaVec& mulElem(const ArmaVec& v)
  {
    for (int i = 0; i < this->dim(); i++)
      ref(i) *= v.get(i);
    return *this;
  }
};

/** Dataset is a set of vectors
 *  Implement n(), dim(), save(), load(), normalize() (text version)
 */
template <typename vec>
class Dataset : public std::vector<vec>
{
public:
  typedef unsigned int uint;
  typedef std::vector<vec> __Base;
  typedef vec vec_type;
public:
  int dim() const { assert(this->size()>0); return __Base::at(0).dim(); }
  int n() const { return __Base::size(); }

  void save(const std::string& fname, arma::file_type type = arma::arma_ascii)
  {
    std::ofstream f(fname.c_str());
    f << this->size() << "\n";
    for (uint i = 0; i < this->size(); i++)
      this->at(i).save(f, type);
    f.close();
  }

  void load(const std::string& fname, arma::file_type type = arma::arma_ascii)
  {
    int n;
    std::string endline;
    std::ifstream f(fname.c_str());
    if (!f.good()) {
      std::cerr << "Bad input file\n";
      return;
    }

    this->clear();
    f >> n;
    std::getline(f, endline);
    for (int i = 0; i < n; i++) {
      this->push_back(vec());
      this->back().load(f, type);
    }
    f.close();
  }

  void print()
  {
    for (uint i = 0; i < this->size(); i++) {
      this->at(i).save(std::cout, arma::arma_ascii);
    }
  }

  void normalize()
  {
    vec mean = this->at(0);
    mean /= n();
    for (int i = 1; i < n(); i++)
      mean += this->at(i) / n();
    for (int i = 1; i < n(); i++)
      this->at(i) -= mean;
    double maxnorm = 0;
    for (int i = 0; i < n(); i++) {
      double norm = arma::norm(this->at(i),2);
      if (maxnorm < norm) maxnorm = norm;
    }
//    for (int i = 0; i < n(); i++)
//      this->at(i) /= maxnorm;
  }
};

/** LabeledDataset is a Dataset with corresponding labels (doubles)
 *  Implement y(), y(int) to access labels, save(), load()
 */
template <typename vec>
class LabeledDataset : public Dataset<vec>
{
public:
  typedef unsigned int              uint;
  typedef Dataset<vec>              __Base;
  typedef std::vector<double>       dvec_type;
  typedef typename __Base::vec_type vec_type;
protected:
  dvec_type y_;
public:
//  int dim() const { assert(this->size()>0); return __Base::at(0).dim(); }
//  int n() const { return __Base::size(); }
  dvec_type& y() { return y_; }
  const dvec_type& y() const { return y_; }
  double y(int i) const { return y_[i]; }
  double& y(int i) { return y_[i]; }

  void save(const std::string& fname, arma::file_type type = arma::arma_ascii)
  {
    std::ofstream f(fname.c_str());
    f << this->size() << "\n";
    for (uint i = 0; i < this->size(); i++) {
      f << y_[i] << "\n";
      this->at(i).save(f, type);
    }
    f.close();
  }

  void load(const std::string& fname, arma::file_type type = arma::arma_ascii)
  {
    int n;
    std::string endline;
    std::ifstream f(fname.c_str());
    if (!f.good()) {
      std::cerr << "Bad input file\n";
      return;
    }
    y_.clear();
    this->clear();
    f >> n;
    std::getline(f, endline);
    for (int i = 0; i < n; i++) {
      y_.push_back(0.0);
      f >> y_.back();
      std::getline(f, endline);
      this->push_back(vec());
      this->back().load(f, type);
    }
    f.close();
  }

  void print()
  {
    for (uint i = 0; i < this->size(); i++) {
      std::cout << y_[i] << "\n";
      this->at(i).save(std::cout, arma::arma_ascii);
    }
  }
};

#endif /* ARMAVEC_H_ */
